import cv2 
from configure import *
import shelve
import numpy as np
import os

class easySample():
    """"
        通过MD5值快速得到各个模型所需要的样本输入
    """
    sample=None
    def __init__(self) -> None:
        pass

    def get_sample(self,name,model):
        if model=="tlsh":
            return self.get_tlsh_sample(name)
        elif model=="siamese_img":
            return self.get_siamese_img_sample(name)
        elif model=="functionSim":
            return self.get_functionSim_sample(name)
        elif model=="MGMN":
            return self.get_MGMN_sample(name)
        elif model=="siamese_graphsage":
            return self.get_siamese_graphsage_sample(name)
        elif model in ["angr","retdec","radare2","ida"]:
            return self.get_disassemble_sample(name,model)
        
    def get_disassemble_sample(self,name,model):
        """
            获得不同工具处理后的样本
            格式一致，只需要换一下路径
        """
        dict={}
        dict["angr"]=r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/angr"
        dict["retdec"]=r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/retdec"
        dict["radare2"]=r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/radare2_functionSim"
        dict["ida"]=r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/functionSim"
        return self.get_functionSim_sample(name,dict[model])


    def get_MGMN_sample(self,name):
        fileroot=functionSim_predata+"//"+name
        res={}
        res["adj"],res["att"]=[],[]
        with shelve.open(fileroot) as file:
            cg=file["cg"]
            cgattr=file["cgattr"]
            caller=file["caller"]#同一行是谁调用了我
            functype=file["funcType"]
            funcs_id=file["func_id"]# name-->ind
        id_funcs={}
        for i in funcs_id.keys():
            id_funcs[funcs_id[i]]=i
        
        #   动态导入函数是[],需要处理一下,
        #   同时生成vtype
        # tempData=[]
        for i in functype.keys():
            # temp=[]
            # if functype[i]=="local":
                # temp=[1,0,0]
            # elif functype[i]=='dynamic import':
            if functype[i]=='dynamic import':
                # temp=[0,1,0]
                cgattr[i]=self.map_api_name_to_eight_bit_vector(id_funcs[i],8,20)
            # else:
                # temp=[0,0,1]
            # tempData.append(temp)
        res["adj"]=np.array(cg)#这里是带边权的
        res["att"]=np.array(cgattr)#应该是8维，但是第一维字符串的提取似乎出现了问题，建议只是用后七维
        # res["vtype"]=np.array(tempData)
        return res

    def get_siamese_graphsage_sample(self,name):
        """
            ***将函数的汇编代码通过asm2vec模型转换成特征向量***
            返回图结构和图节点的汇编代码
            
        """
        radara2SamplePath="/home/cyw/projects/function_sim_project/all_data/sampleDatas/radare2_predata/"
        
        fileroot=radara2SamplePath+name
        res={}
        res["adj"],res["att"]=[],[]
        # 增加判断文件是否存在，不然shelve会生成一个默认的
        if os.path.exists(fileroot+".dir"):
            with shelve.open(fileroot) as file:
                sampleInf=file["data"]
                # size:样本中函数的个数
                # discode:函数名及其对应的反汇编代码
                # adj:邻接矩阵
                # calls：函数间调用关系，都是函数名
                # name_to_id:函数名与下标的映射
                res["adj"]=sampleInf["adj"]
                res["name_to_id"]=sampleInf["name_to_id"]
                res["adj"]=sampleInf["adj"]
                res["att"]=sampleInf["att"] 
        else:
            raise ValueError("样本embedding不存在!!!")
        return res
        
    def get_functionSim_sample(self,name,filePath=r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/functionSim"):
        """
            路径输入默认是ida进行反汇编
            输出：
                adj 邻接矩阵 b*n*n 
                att 特征矩阵 b*n*d
                vtype 类型矩阵 b*n*3  3种类型
        """
        fileroot=filePath+"//"+name
        res={}
        res["adj"],res["att"],res["vtype"]=[],[],[]
        with shelve.open(fileroot) as file:
            cg=file["cg"]
            cgattr=file["cgattr"]
            # caller=file["caller"]#同一行是谁调用了我
            # 这里的名称似乎写错了。。。
            try:
                functype=file["funcType"]
            except Exception as e:
                functype=file["functype"]
            try:
                funcs_id=file["func_id"]# name-->ind
            except Exception as e:
                funcs_id=file["funcs_id"]# name-->ind
        id_funcs={}
        for i in funcs_id.keys():
            id_funcs[funcs_id[i]]=i

        #   动态导入函数是[],需要处理一下,
        #   同时生成vtype
        tempData=[]
        for i in functype.keys():
            temp=[]
            if functype[i]=="local":
                temp=[1,0,0]
            elif functype[i]=='dynamic import':
                temp=[0,1,0]
                cgattr[i]=self.map_api_name_to_eight_bit_vector(id_funcs[i],8,20)
            else:
                temp=[0,0,1]
            tempData.append(temp)
        res["adj"]=np.array(cg)#这里是带边权的
        res["att"]=np.array(cgattr)#应该是8维，但是第一维字符串的提取似乎出现了问题，建议只是用后七维
        res["vtype"]=np.array(tempData)
        return res
        
    def get_tlsh_sample(self,name):
        return name

    def get_siamese_img_sample(self,name):
        imgPath = "/home/cyw/projects/function_sim_project/all_data/sampleDatas/img"
        res=cv2.imread(imgPath+"/"+name+".png",cv2.IMREAD_UNCHANGED)
        return res
    
    def map_api_name_to_eight_bit_vector(self,function_name,embeddingSize,numLim):
        """
            将api转换成一个embeddingSize位的向量，每一位使用numLim取余
        """
        function_bytes = function_name.encode('utf-8')
        eight_bit_vector = [0] * embeddingSize
        for i, byte in enumerate(function_bytes):
            index = i % embeddingSize
            eight_bit_vector[index] += byte
            eight_bit_vector[index] %= numLim
        return eight_bit_vector

if __name__=="__main__":    
    a=easySample()
    b=a.get_sample("ff14cb827538ea4e65dcdc26e5d6854b","functionSim")

    print(b)